/* 
 *   This file is part of OpenTLD.
 *
 *   OpenTLD is free software: you can redistribute it and/or modify
 *   it under the terms of the GNU General Public License as published by
 *    the Free Software Foundation, either version 3 of the License, or
 *   (at your option) any later version.
 *
 *   OpenTLD is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *   GNU General Public License for more details.
 *
 *   You should have received a copy of the GNU General Public License
 *   along with OpenTLD.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#include "TLD.h"

#include "NNClassifier.h"
#include "TLDUtil.h"

#include <iostream>
#include <highgui.h>

using namespace std;
using namespace cv;

namespace tld
{

TLD::TLD() {
	trackerEnabled = true;
	detectorEnabled = true;
	learningEnabled = true;
	alternating = false;
	valid = false;
	wasValid = false;
	learning = false;
	currBB = prevBB = NULL;

	detectorCascade = new DetectorCascade();
	nnClassifier = detectorCascade->nnClassifier;
	medianFlowTracker = new MedianFlowTracker();

	overlap_pth = 0.6;
	overlap_nth = 0.2;
	thetaFP_fuse = 0.5;
	thetaTP_fuse = 0.65;
}

TLD::~TLD() {
	storeCurrentData();
	
	delete detectorCascade;
	delete medianFlowTracker;
}

void TLD::release() {
	detectorCascade->release();
	medianFlowTracker->cleanPreviousData();
	delete currBB;
}

void TLD::storeCurrentData() {
	prevImg.release();
	prevImg = currImg; //Store old image (if any)
	delete prevBB;
	prevBB = currBB;        //Store old bounding box (if any)
	
	detectorCascade->cleanPreviousData(); //Reset detector results
	medianFlowTracker->cleanPreviousData();

	wasValid = valid;
}

void TLD::selectObject(const Mat &img, Rect *bb) {
	//Delete old object
	detectorCascade->release();

	detectorCascade->objWidth = bb->width;
	detectorCascade->objHeight = bb->height;

	//Init detector cascade
	detectorCascade->init();
	
	currImg = img;
	currBB = tldCopyRect(bb);
	currConf = 1;
	valid = true;
	initialLearning();
}

void TLD::processImage(const Mat &img, bool showTrackResult) {
	cout << "TLD Processing...\n";
	storeCurrentData();
	Mat grey_frame;
	cvtColor(img, grey_frame, CV_RGB2GRAY);
	currImg = grey_frame; // Store new image , right after storeCurrentData();
	
	if(trackerEnabled) {
		medianFlowTracker->track(prevImg, currImg, prevBB, showTrackResult);
	}

	if(detectorEnabled && (!alternating || medianFlowTracker->trackerBB == NULL)) {
		detectorCascade->detect(grey_frame);
	}
	
	fuseHypotheses();
	
	learn();

}

void TLD::fuseHypotheses() {
	Rect *trackerBB = medianFlowTracker->trackerBB;
	int numClusters = detectorCascade->detectionResult->numClusters;
	Rect *detectorBB = detectorCascade->detectionResult->detectorBB;

	currBB = NULL;
	currConf = 0;
	valid = false;

	float confDetector = 0;
	
	if(numClusters == 1) {
		confDetector = nnClassifier->classifyBB(currImg, detectorBB);
	}

	if(trackerBB != NULL) {
		float confTracker = nnClassifier->classifyBB(currImg, trackerBB);
		if(numClusters == 1 && confDetector > confTracker && tldOverlapRectRect(*trackerBB, *detectorBB) < 0.5) {
			currBB = tldCopyRect(detectorBB);
			currConf = confDetector;
		} else {
			currBB = tldCopyRect(trackerBB);
			currConf = confTracker;
			
			if(confTracker > thetaTP_fuse) {
				valid = true;
			} else if(wasValid && confTracker > thetaFP_fuse) {
				valid = true;
			}
		}
	} else if(numClusters == 1) {
		currBB = tldCopyRect(detectorBB);
		currConf = confDetector;
	}
}

void TLD::initialLearning() {
    	cout << "TLD Initial Learning...\n";
	learning = true; // This is just for display purposes
	
	DetectionResult *detectionResult = detectorCascade->detectionResult;
	detectorCascade->detect(currImg);

	//This is the positive patch
	NormalizedPatch patch;
	tldExtractNormalizedPatchRect(currImg, currBB, patch.values);
	patch.positive = 1;

	float initVar = tldCalcVariance(patch.values, TLD_PATCH_SIZE * TLD_PATCH_SIZE);
	detectorCascade->varianceFilter->minVar = initVar / 2; // It's ok to change to other reasonable factor
	
	float *overlap = new float[detectorCascade->numWindows];
	tldOverlapRect(detectorCascade->windows, detectorCascade->numWindows, currBB, overlap);
	
	//Add all bounding boxes with high overlap
	vector< pair<int, float> > positiveIndices;
	vector<int> negativeIndices;

	//First: Find overlapping positive and negative patches
	for(int i = 0; i < detectorCascade->numWindows; i++) {
		if(overlap[i] > overlap_pth) {
			positiveIndices.push_back(pair<int, float>(i, overlap[i]));
		}
		if(overlap[i] < overlap_nth) {
			float variance = detectionResult->variances[i];
			if(!detectorCascade->varianceFilter->enabled 
			|| variance > detectorCascade->varianceFilter->minVar) {  //TODO: This check is unnecessary if minVar would be set before calling detect. 
				negativeIndices.push_back(i);
			}
		}
	}

	sort(positiveIndices.begin(), positiveIndices.end(), tldSortByOverlapDesc);
	vector<NormalizedPatch> patches;
	patches.push_back(patch); //Add first patch to patch list
	
	int numIterations = std::min<size_t>(positiveIndices.size(), 10); //Take at most 10 bounding boxes (sorted by overlap)

	for(int i = 0; i < numIterations; i++) {
		int idx = positiveIndices.at(i).first;
		//Learn this bounding box
		//TODO: Somewhere here image warping might be possible
		detectorCascade->ensembleClassifier->learn(&detectorCascade->windows[TLD_WINDOW_SIZE * idx], true, &detectionResult->featureVectors[detectorCascade->numTrees * idx]);
	}

	srand(1); 
	random_shuffle(negativeIndices.begin(), negativeIndices.end());

	//Choose 100 random patches for negative examples
	for(size_t i = 0; i < std::min<size_t>(100, negativeIndices.size()); i++) {
		int idx = negativeIndices.at(i);
		NormalizedPatch patch;
		tldExtractNormalizedPatchBB(currImg, &detectorCascade->windows[TLD_WINDOW_SIZE * idx], patch.values);
		patch.positive = 0;
		patches.push_back(patch);
	}
	
	detectorCascade->nnClassifier->learn(patches);
	delete[] overlap;
}

//Do this when current trajectory is valid
void TLD::learn() {
	if(!learningEnabled || !valid || !detectorEnabled) {
		learning = false;
		return;
	}

	learning = true;
	DetectionResult *detectionResult = detectorCascade->detectionResult;

	if(!detectionResult->containsValidData) {
		detectorCascade->detect(currImg);
	}
	
	//This is the positive patch
	NormalizedPatch patch;
	tldExtractNormalizedPatchRect(currImg, currBB, patch.values);

	float *overlap = new float[detectorCascade->numWindows];
	tldOverlapRect(detectorCascade->windows, detectorCascade->numWindows, currBB, overlap);

	//Add all bounding boxes with high overlap

	vector<pair<int, float> > positiveIndices;
	vector<int> negativeIndices;
	vector<int> negativeIndicesForNN;

	//First: Find overlapping positive and negative patches
	for(int i = 0; i < detectorCascade->numWindows; i++) {
		if(overlap[i] > overlap_pth) {
			positiveIndices.push_back(pair<int, float>(i, overlap[i]));
		}
		if(overlap[i] < overlap_nth) {
			if(!detectorCascade->ensembleClassifier->enabled || detectionResult->posteriors[i] > 0.5) {
				negativeIndices.push_back(i);
			}
			if(!detectorCascade->ensembleClassifier->enabled || detectionResult->posteriors[i] > 0.5) {
				negativeIndicesForNN.push_back(i);
			}
		}
	}
	
	sort(positiveIndices.begin(), positiveIndices.end(), tldSortByOverlapDesc);

	cout << "......positive size of ensemble classifier: " << positiveIndices.size() << "\n";
	cout << "......negative size of ensemble classifier: " << negativeIndices.size() << "\n";
	cout << "......negative size of NN classifier: " << negativeIndicesForNN.size() << "\n";

	vector<NormalizedPatch> patches;
	
	patch.positive = 1;
	patches.push_back(patch);

	int numIterations = std::min<size_t>(positiveIndices.size(), 10); //Take at most 10 bounding boxes (sorted by overlap)

	for(size_t i = 0; i < negativeIndices.size(); i++) {
        	int idx = negativeIndices.at(i);
		//TODO: Somewhere here image warping might be possible
		detectorCascade->ensembleClassifier->learn(&detectorCascade->windows[TLD_WINDOW_SIZE * idx], false, &detectionResult->featureVectors[detectorCascade->numTrees * idx]);
	}
	
	//TODO: Randomization might be a good idea
	for(int i = 0; i < numIterations; i++) {
		int idx = positiveIndices.at(i).first;
		//TODO: Somewhere here image warping might be possible
		detectorCascade->ensembleClassifier->learn(&detectorCascade->windows[TLD_WINDOW_SIZE * idx], true, &detectionResult->featureVectors[detectorCascade->numTrees * idx]);
	}
	
	for(size_t i = 0; i < negativeIndicesForNN.size(); i++) {
		int idx = negativeIndicesForNN.at(i);

		NormalizedPatch patch;
		tldExtractNormalizedPatchBB(currImg, &detectorCascade->windows[TLD_WINDOW_SIZE * idx], patch.values);
		patch.positive = 0;
		patches.push_back(patch);
	}
	
	detectorCascade->nnClassifier->learn(patches);
	
	//cout << "NN has now " << detectorCascade->nnClassifier->truePositives->size() << " positives and " 
	//	<< detectorCascade->nnClassifier->falsePositives->size() << " negatives.\n";

	delete[] overlap;
}

typedef struct
{
    int index;
    int P;
    int N;
} TldExportEntry;



} /* namespace tld */
